#include "PaddleOcrOpenvino.h"

using namespace std;
using namespace cv;

#define PI (3.1415926)

// extern std::shared_ptr<spdlog::logger> kylog;

PaddleOcrOpenvino::PaddleOcrOpenvino(void) 
{

}

PaddleOcrOpenvino::~PaddleOcrOpenvino() 
{
    
}

bool PaddleOcrOpenvino::Init(char *weight, char *dictFile)
{
    bool bret = false;

    model = core.read_model(weight);

    compiled_model = core.compile_model(model, "CPU");
    infer_request = compiled_model.create_infer_request();

    input_tensor = infer_request.get_input_tensor();
    tensor_shape = input_tensor.get_shape();

    inChannel = tensor_shape[1];
    inHeight = tensor_shape[2];
    inWidth = tensor_shape[3];

    bret = InitDict(dictFile, dict);
    if(!bret)
    {
        // kylog->info("open dictFile {} failed!", dictFile);
    }

    isInitOk = true;
    bret = true;

    return bret;
}

bool PaddleOcrOpenvino::InitDict(char *file, vector<string> &dict)
{
    bool bret = false;

    ifstream readFile;
    readFile.open(file);

    dict.push_back("-");

    if (readFile)
    {
        string temp;
        while (getline(readFile, temp))
        {
            dict.push_back(temp);
        }
        readFile.close();
        bret = true;

        // for(auto i : dict)
        // {
        //     cout << i << endl;
        // }
    }

    return bret;
}

bool PaddleOcrOpenvino::Ocr2rd(cv::Mat &img, cv::Rect &rect, string &result, int &iscore)
{
    int temp = 0;
    return Ocr2rd(img, rect, result, iscore, temp);
}

bool PaddleOcrOpenvino::Ocr2rd(cv::Mat &img, cv::Rect &rect, std::string &result, int &iscore, int &lastScore)
{
    bool bret = false;
    cv::Rect rectFix;

    if(rect.x <= 0 || rect.y <= 0 || rect.x + rect.width >= img.cols || rect.y + rect.height >= img.rows
            || rect.width <= 0 || rect.height <= 0)
    {
        // kylog->info("rect out of boarder!");
        return bret;
    }

    memcpy(&rectFix, &rect, sizeof(rectFix));
    AdjectRect(img, rectFix, 3, 1, 8, 12);
    // AdjectRect(img, rectFix, 1, 1, 2, 2);

    cv::Mat imageParkid = img(rectFix);
    bret = OcrPaddleProcess(imageParkid, result, iscore, lastScore);

    rect.x = rectFix.x;
    rect.y = rectFix.y;
    rect.width = rectFix.width;
    rect.height = rectFix.height;

    return bret;
}

bool PaddleOcrOpenvino::Ocr(cv::Mat &img, cv::Rect &rect, std::string &result, int &iscore, int &lastScore)
{
    bool bret = false;

    if(rect.x <= 0 || rect.y <= 0 || rect.x + rect.width >= img.cols || rect.y + rect.height >= img.rows
            || rect.width <= 0 || rect.height <= 0)
    {
        // kylog->info("rect out of boarder!");
        return bret;
    }

    cv::Mat imageParkid = img(rect);

    return OcrPaddleProcess(imageParkid, result, iscore, lastScore);
}

bool PaddleOcrOpenvino::Ocr(cv::Mat &img, cv::Rect &rect, string &result, int &iscore)
{
    bool bret = false;

    if(rect.x <= 0 || rect.y <= 0 || rect.x + rect.width >= img.cols || rect.y + rect.height >= img.rows
            || rect.width <= 0 || rect.height <= 0)
    {
        // kylog->info("rect out of boarder!");
        return bret;
    }

    cv::Mat imageParkid = img(rect);
    bret = OcrPaddleProcess(imageParkid, result, iscore);

    return bret;
}

bool PaddleOcrOpenvino::Ocr(cv::Mat &img, std::string &result, int &iscore)
{
    return OcrPaddleProcess(img, result, iscore);
}

bool PaddleOcrOpenvino::OcrPaddleProcess(cv::Mat &img, std::string &result, int &iscore)
{
    int temp = 0;
    return OcrPaddleProcess(img, result, iscore, temp);
}

bool PaddleOcrOpenvino::OcrPaddleProcess(cv::Mat &img, std::string &result, int &iscore, int &lastScore)
{
    bool bret = false;
    Mat blob_image;

    paddleResize(img, blob_image);
    normalize(&blob_image, mean, std, true);

    float* image_data = (float*)input_tensor.data();
    size_t imageSize = inWidth * inHeight;

    for (size_t pid = 0; pid < imageSize; ++pid)
    {
        for (size_t ch = 0; ch < inChannel; ++ch)
        {
            image_data[imageSize*ch + pid] = (blob_image.at<cv::Vec3f>(pid)[ch]);
        }
    }

    infer_request.infer();

    auto output_tensor = infer_request.get_output_tensor();
    float* detection = (float*)output_tensor.data();
    
    auto output_shape = output_tensor.get_shape();

    int outputCount = output_shape[0];
    int outputNum = output_shape[1];
    int outputLen = output_shape[2];

    for (int m = 0; m < outputCount; m++)
    {
        int argmax_idx;
        int last_index = 0;
        float score = 0.f;
        int count = 0;
        float max_value = 0.0f;

        for (int n = 0; n < outputNum; n++)
        {
            // get idx
            argmax_idx = (max_element(
                &detection[(m * outputNum + n) * outputLen],
                &detection[(m * outputNum + n + 1) * outputLen])) - &detection[(m * outputNum + n) * outputLen];
            // get score
            max_value = float(*std::max_element(
                &detection[(m * outputNum + n) * outputLen],
                &detection[(m * outputNum + n + 1) * outputLen]));

            if (argmax_idx > 0 && (!(n > 0 && argmax_idx == last_index)))
            {
                if(argmax_idx < dict.size())
                {
                    score += max_value;
                    count += 1;
                    result += dict[argmax_idx];
                    lastScore = max_value * 100;
                }
            }
            last_index = argmax_idx;
        }
        score /= count;
        if (std::isnan(score))
        {
            continue;
        }
        iscore = score * 100;
        // kylog->info("result: {}, score = {}", result.c_str(), iscore);
        // printf("%s score = %f\r\n", result.c_str(), score);
    }

    bret = true;

    return bret;
}

void PaddleOcrOpenvino::paddleResize(cv::Mat &img, cv::Mat &resize_img)
{
    int imgH = inHeight;
    int imgW = inWidth;
    float max_wh_ratio = imgW * 1.0 / imgH;

    int h = img.rows;
    int w = img.cols;
    float wh_ratio = w * 1.0 / h;
    max_wh_ratio = max(max_wh_ratio, wh_ratio);

    int imgC;
    imgC = inChannel;
    imgH = inHeight;
    imgW = inWidth;

    imgW = int(imgH * max_wh_ratio);

    float ratio = float(img.cols) / float(img.rows);
    int resize_w, resize_h;

    if (ceilf(imgH * ratio) > imgW)
        resize_w = imgW;
    else
        resize_w = int(ceilf(imgH * ratio));

    cv::resize(img, resize_img, cv::Size(resize_w, imgH), 0.f, 0.f,
            cv::INTER_LINEAR);
    cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0,
            int(imgW - resize_img.cols), cv::BORDER_CONSTANT,
            {127, 127, 127});  
}

void PaddleOcrOpenvino::normalize(cv::Mat *im, float *mean,
        float *scale, bool is_scale)
{
    double e = 1.0;
    if (is_scale)
    {
        e /= 255.0;
    }
    (*im).convertTo(*im, CV_32FC3, e);
    std::vector<cv::Mat> bgr_channels(3);
    cv::split(*im, bgr_channels);
    for (auto i = 0; i < bgr_channels.size(); i++)
    {
        bgr_channels[i].convertTo(bgr_channels[i], CV_32FC1, 1.0 * scale[i],
                                  (0.0 - mean[i]) * scale[i]);
    }
    cv::merge(bgr_channels, *im);
}

void PaddleOcrOpenvino::AdjectRect(cv::Mat &oriMat, cv::Rect &rect, int x0, int y0, int w0, int h0)
{
    int x, y, w, h;
    x = rect.x;
    y = rect.y;
    w = rect.width;
    h = rect.height;

    if(x - x0 >= 0)
    {
        x = x - x0;
    }
    else 
    {
        x = 0;
    }
    if(y - y0 >= 0)
    {
        y = y - y0;//can no fix
    }
    else 
    {
        y = 0;
    }
    if(x + w + w0 <= oriMat.cols - 1)
    {
        w = w + w0;
    }
    else
    {
        w = oriMat.cols - 1 - x;
    }
    if(y + h + h0 <= oriMat.rows - 1)
    {
        h = h + h0;//can add to 8~10
    }
    else
    {
        h = oriMat.rows - 1 - y;
    }

    rect.x = x;
    rect.y = y;
    rect.width = w;
    rect.height = h;
}

void PaddleOcrOpenvino::edlinesProc(cv::Mat &imageParkid, MLine &line1, MLine &line2)
{
    double max1 = 0, max2 = 0, angle1 = 0, angle2 = 0, k1 = 0, k2 = 0, y1 = 0, h = 0, max1LineX = 0, max1LineY = 0;
    struct LS maxLine1(cv::Point2d(0, 0), cv::Point2d(0, 0)), maxLine2(cv::Point2d(0, 0), cv::Point2d(0, 0));
    bool isLine1Ok = false;
    bool isLine2Ok = false;

    memset(&line1, 0, sizeof(line1));
    memset(&line2, 0, sizeof(line2));

    cv::Mat image;
    cv::cvtColor(imageParkid, image, cv::COLOR_RGB2GRAY);

    EDLines lineHandler = EDLines(image);

    std::vector<LS> lines = lineHandler.getLines();
    for(auto line : lines)
    {
        double dis = sqrt((line.start.x - line.end.x) * (line.start.x - line.end.x) + (line.start.y - line.end.y) * (line.start.y - line.end.y));
        if(dis > max1 && line.start.x != line.end.x)
        {
            max1 = dis;
            memcpy(&maxLine1, &line, sizeof(maxLine1));
            angle1 = atan(1.0 * (line.start.y - line.end.y)/(line.start.x - line.end.x));
            k1 = 1.0 * (line.start.y - line.end.y)/(line.start.x - line.end.x);
            y1 = k1 * image.cols - k1 * line.start.x + line.start.y;
            max1LineX = line.end.x;
            max1LineY = line.end.y;
            isLine1Ok = true;
        }
    }
    if(isLine1Ok)
    {
        line1.x0 = 0;
        line1.x1 = imageParkid.cols;
        line1.y0 = -k1 * maxLine1.start.x + maxLine1.start.y;
        line1.y1 = k1 * (line1.x1 - maxLine1.start.x) + maxLine1.start.y;
        line1.angle = angle1;
    }

    for(auto line : lines)
    {
        double dis = sqrt((line.start.x - line.end.x) * (line.start.x - line.end.x) + (line.start.y - line.end.y) * (line.start.y - line.end.y));
        double angle = 0.0;
        double y1, y2;
        double tempk;

        if(line.start.x == line.end.x)
        {
            angle = PI / 2;
            tempk = 999999.9;
            y2 = tempk * image.cols - tempk * line.start.x + line.start.y;
        }
        else 
        {
            angle = atan(1.0 * (line.start.y - line.end.y)/(line.start.x - line.end.x));
            tempk = 1.0 * (line.start.y - line.end.y)/(line.start.x - line.end.x);
            y2 = tempk * image.cols - tempk * line.start.x + line.start.y;
        }

        h = fabs(k1*line.end.x - line.end.y - k1*max1LineX + max1LineY) / sqrt(k1*k1 + 1);

        if(dis > max2 && dis < max1 && fabs(angle - angle1) < PI / 36 && h > 15)
        {
            max2 = dis;
            memcpy(&maxLine2, &line, sizeof(maxLine2));
            isLine2Ok = true;
            angle2 = atan(1.0 * (line.start.y - line.end.y)/(line.start.x - line.end.x));
            k2 = tempk;
        }
    }
    if(isLine2Ok)
    {
        line2.x0 = 0;
        line2.x1 = imageParkid.cols;
        line2.y0 = -k2 * maxLine2.start.x + maxLine2.start.y;
        line2.y1 = k2 * (line2.x1 - maxLine2.start.x) + maxLine2.start.y;
        line2.angle = angle2;
        // printf("%lf, %lf\r\n", Hudu2Jiaodu(angle1), Hudu2Jiaodu(angle2));       
    }
}

bool PaddleOcrOpenvino::parkidProcess(cv::Mat &oriMat, cv::Mat &retMat, cv::Rect &box)
{
    bool bret = false;
    cv::Rect rectFix, parkRect;
    MLine line1, line2;

    if(box.x <= 0 || box.y <= 0 || box.x + box.width >= oriMat.cols || box.y + box.height >= oriMat.rows
            || box.width <= 0 || box.height <= 0)
    {
        // kylog->info("rect out of boarder!");
        return bret;
    }

    memcpy(&rectFix, &box, sizeof(rectFix));

    // AdjectRect(oriMat, rectFix, 3, 1, 8, 12);
    cv::Mat imageParkid = oriMat(rectFix);
    edlinesProc(imageParkid, line1, line2);
    
    cv::Mat imshowImg = imageParkid.clone();
    cv::line(imshowImg, Point(line1.x0, line1.y0), Point(line1.x1, line1.y1), cv::Scalar(32,0,200), 3, 8);
    cv::line(imshowImg, Point(line2.x0, line2.y0), Point(line2.x1, line2.y1), cv::Scalar(200,0,32), 3, 8);
    imshow("parkid1", imshowImg);
    // cv::waitKey(0);

    /** line1 line2 ok **/
    if(line1.y1 > 0 && line2.y1 > 0)
    {
        parkRect.x = box.x;
        parkRect.width = imageParkid.cols;
        if(line1.y0 > line2.y0)
        {
            parkRect.y = line2.y0 > line2.y1 ? line2.y1 : line2.y0;
            parkRect.height = line1.y0 > line1.y1 ? line1.y0 : line1.y1;
            parkRect.height = parkRect.height - parkRect.y;
            parkRect.y += box.y;
        }
        else 
        {
            parkRect.y = line1.y0 > line1.y1 ? line1.y1 : line1.y0;
            parkRect.height = line2.y0 > line2.y1 ? line2.y0 : line2.y1;
            parkRect.height = parkRect.height - parkRect.y;
            parkRect.y += box.y;
        }

        AdjectRect(oriMat, parkRect, 1, 1, 2, 2);
        cv::Mat ori = oriMat(parkRect);
        cv::Mat dst;

        //imshow("ori", ori);

        double jiaodu = Hudu2Jiaodu(line1.angle);
        // printf("jiaodu = %lf\r\n", jiaodu);

        if(fabs(line1.angle) > PI / 36)
        {
            cv::Point2f center(ori.cols / 2, ori.rows / 2); //定义旋转中心

            cv::Mat M = cv::getRotationMatrix2D(center, jiaodu, 1.0); //定义旋转矩阵
            cv::warpAffine(ori, dst, M, ori.size()); //进行旋转操作，输出结果保存在dst Mat对象中

            int maxh = abs(line1.y0 - line2.y0) > abs(line1.y1 - line2.y1) ? abs(line1.y0 - line2.y0) : abs(line1.y1 - line2.y1);

            cv::Rect fixRect;
            fixRect.x = 0;
            fixRect.y = ori.rows / 2 - maxh / 2;
            fixRect.width = ori.cols;
            fixRect.height = maxh;

            AdjectRect(dst, fixRect, 2, 5, 4, 10);

            dst = dst(fixRect);
        }
        else 
        {
            dst = ori;
        }

        imshow("parkid", dst);
        // cv::waitKey(0);
        retMat = dst.clone();
        bret = true;
    }

    return bret;
}

bool PaddleOcrOpenvino::OcrEDLine(cv::Mat &img, cv::Mat &procImg, cv::Rect &box, std::string &result, int &iscore)
{
    int temp = 0;
    return OcrEDLine(img, procImg, box, result, iscore, temp);
}

bool PaddleOcrOpenvino::OcrEDLine(cv::Mat &img, cv::Mat &procImg, cv::Rect &box, std::string &result, int &iscore, int &lastScore)
{
    bool bret = false;
    cv::Mat detMat;

    bret = parkidProcess(img, detMat, box);
    if(bret)
    {
        bret = OcrPaddleProcess(detMat, result, iscore, lastScore);
        if(bret)
        {
            procImg = detMat.clone();
        }
    }

    return bret;
}

double PaddleOcrOpenvino::Hudu2Jiaodu(double hudu)
{
    double ret;

    ret = hudu * 180.0 / PI;

    return ret;
}
